package middleware

import (
	"context"
	"github.com/gin-gonic/gin"
	"github.com/go-redis/redis/v8"
	"net/http"
	"strconv"
)

// RateLimit 中间件：基于 Redis 的令牌桶限流
func RateLimit(client *redis.Client, limit int) gin.HandlerFunc {
	return func(c *gin.Context) {
		apiKey := c.GetHeader("X-API-Key")
		key := "rate_limit:" + apiKey

		// 当前时间戳（秒）
		window := 60 // 每分钟限流

		// 使用 Redis 的 Lua 脚本实现原子操作
		script := `
            local current = redis.call("GET", KEYS[1])
            if current and tonumber(current) > tonumber(ARGV[1]) then
                return false
            end
            redis.call("SET", KEYS[1], tonumber(current or 0) + 1, "EX", ARGV[2])
            return true
        `

		allowed, err := client.Eval(context.Background(), script, []string{key}, strconv.Itoa(limit), strconv.Itoa(window)).Bool()
		if err != nil || !allowed {
			c.AbortWithStatusJSON(http.StatusTooManyRequests, gin.H{"error": "rate limit exceeded"})
			return
		}

		c.Next()
	}
}
